#!/usr/bin/env python3
"""
Hyperliquid API Client

This module provides a comprehensive API client for interacting with
the Hyperliquid exchange, supporting both spot and perpetual trading.
"""

import asyncio
import hashlib
import hmac
import json
import logging
import time
from decimal import Decimal
from typing import Dict, List, Optional, Any, Union

import aiohttp
import websockets
from dataclasses import dataclass


@dataclass
class OrderBook:
    """Order book data structure"""
    symbol: str
    bids: List[List[str]]  # [price, size]
    asks: List[List[str]]  # [price, size]
    timestamp: int


@dataclass
class Position:
    """Position data structure"""
    symbol: str
    side: str  # "long" or "short"
    size: Decimal
    entry_price: Decimal
    mark_price: Decimal
    unrealized_pnl: Decimal
    margin: Decimal


@dataclass
class Order:
    """Order data structure"""
    order_id: str
    symbol: str
    side: str  # "buy" or "sell"
    order_type: str  # "market", "limit", etc.
    amount: Decimal
    price: Optional[Decimal]
    status: str
    filled_amount: Decimal
    timestamp: int


@dataclass
class Trade:
    """Trade data structure"""
    trade_id: str
    symbol: str
    side: str
    amount: Decimal
    price: Decimal
    timestamp: int
    fee: Decimal


class HyperliquidAPIClient:
    """
    Hyperliquid API Client for both spot and perpetual trading
    """
    
    def __init__(self, 
                 api_key: Optional[str] = None,
                 api_secret: Optional[str] = None,
                 testnet: bool = True,
                 base_url: Optional[str] = None,
                 ws_url: Optional[str] = None):
        """
        Initialize the Hyperliquid API client
        
        Args:
            api_key: API key for authentication
            api_secret: API secret for authentication
            testnet: Whether to use testnet
            base_url: Custom base URL
            ws_url: Custom WebSocket URL
        """
        self.api_key = api_key
        self.api_secret = api_secret
        self.testnet = testnet
        
        # Set URLs based on testnet flag
        if testnet:
            self.base_url = base_url or "https://api.hyperliquid-testnet.xyz"
            self.ws_url = ws_url or "wss://api.hyperliquid-testnet.xyz/ws"
        else:
            self.base_url = base_url or "https://api.hyperliquid.xyz"
            self.ws_url = ws_url or "wss://api.hyperliquid.xyz/ws"
        
        # Session and WebSocket
        self._session: Optional[aiohttp.ClientSession] = None
        self._ws: Optional[websockets.WebSocketServerProtocol] = None
        
        # Logger
        self._logger = logging.getLogger(__name__)
        
        # Rate limiting
        self._last_request_time = 0
        self._min_request_interval = 0.1  # 100ms between requests
        
        self._logger.info(f"Initialized Hyperliquid API client (testnet={testnet})")
    
    async def __aenter__(self):
        """Async context manager entry"""
        await self.connect()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit"""
        await self.disconnect()
    
    async def connect(self):
        """Initialize HTTP session"""
        if not self._session:
            timeout = aiohttp.ClientTimeout(total=30)
            self._session = aiohttp.ClientSession(timeout=timeout)
            self._logger.info("HTTP session initialized")
    
    async def disconnect(self):
        """Close HTTP session and WebSocket"""
        if self._session:
            await self._session.close()
            self._session = None
            self._logger.info("HTTP session closed")
        
        if self._ws:
            await self._ws.close()
            self._ws = None
            self._logger.info("WebSocket connection closed")
    
    def _generate_signature(self, timestamp: str, method: str, path: str, body: str = "") -> str:
        """Generate HMAC signature for authentication"""
        if not self.api_secret:
            raise ValueError("API secret required for authentication")
        
        message = f"{timestamp}{method}{path}{body}"
        signature = hmac.new(
            self.api_secret.encode('utf-8'),
            message.encode('utf-8'),
            hashlib.sha256
        ).hexdigest()
        
        return signature
    
    async def _rate_limit(self):
        """Implement rate limiting"""
        current_time = time.time()
        time_since_last = current_time - self._last_request_time
        
        if time_since_last < self._min_request_interval:
            await asyncio.sleep(self._min_request_interval - time_since_last)
        
        self._last_request_time = time.time()
    
    async def _make_request(self, 
                           method: str, 
                           endpoint: str, 
                           params: Optional[Dict] = None,
                           data: Optional[Dict] = None,
                           authenticated: bool = False) -> Dict[str, Any]:
        """Make HTTP request to Hyperliquid API"""
        if not self._session:
            await self.connect()
        
        await self._rate_limit()
        
        url = f"{self.base_url}{endpoint}"
        headers = {
            "Content-Type": "application/json",
            "User-Agent": "HyperliquidBot/1.0"
        }
        
        # Add authentication headers if required
        if authenticated and self.api_key and self.api_secret:
            timestamp = str(int(time.time() * 1000))
            body = json.dumps(data) if data else ""
            signature = self._generate_signature(timestamp, method, endpoint, body)
            
            headers.update({
                "HX-ACCESS-KEY": self.api_key,
                "HX-ACCESS-SIGN": signature,
                "HX-ACCESS-TIMESTAMP": timestamp
            })
        
        try:
            async with self._session.request(
                method=method,
                url=url,
                params=params,
                json=data,
                headers=headers
            ) as response:
                response_data = await response.json()
                
                if response.status != 200:
                    self._logger.error(f"API request failed: {response.status} - {response_data}")
                    raise Exception(f"API request failed: {response.status} - {response_data}")
                
                return response_data
                
        except Exception as e:
            self._logger.error(f"Request error: {e}")
            raise
    
    # Market Data Methods
    
    async def get_exchange_info(self) -> Dict[str, Any]:
        """Get exchange information"""
        return await self._make_request("GET", "/info")
    
    async def get_order_book(self, symbol: str, limit: int = 100) -> OrderBook:
        """Get order book for a symbol"""
        params = {"coin": symbol, "nSigFigs": limit}
        data = await self._make_request("POST", "/info", data={"type": "l2Book", **params})
        
        return OrderBook(
            symbol=symbol,
            bids=data.get("levels", [[]])[0] if data.get("levels") else [],
            asks=data.get("levels", [[], []])[1] if data.get("levels") else [],
            timestamp=int(time.time() * 1000)
        )
    
    async def get_ticker(self, symbol: str) -> Dict[str, Any]:
        """Get ticker information for a symbol"""
        data = await self._make_request("POST", "/info", data={"type": "allMids"})
        return data.get(symbol, {})
    
    async def get_funding_rate(self, symbol: str) -> Dict[str, Any]:
        """Get funding rate for a perpetual symbol"""
        data = await self._make_request("POST", "/info", data={"type": "fundingHistory", "coin": symbol})
        return data
    
    async def get_klines(self, 
                        symbol: str, 
                        interval: str = "1m", 
                        start_time: Optional[int] = None,
                        end_time: Optional[int] = None,
                        limit: int = 500) -> List[List[str]]:
        """Get kline/candlestick data"""
        params = {
            "coin": symbol,
            "interval": interval,
            "startTime": start_time,
            "endTime": end_time
        }
        
        # Remove None values
        params = {k: v for k, v in params.items() if v is not None}
        
        data = await self._make_request("POST", "/info", data={"type": "candleSnapshot", **params})
        return data.get("candles", [])
    
    # Account Methods
    
    async def get_account_info(self) -> Dict[str, Any]:
        """Get account information"""
        return await self._make_request("POST", "/info", 
                                      data={"type": "clearinghouseState", "user": self.api_key},
                                      authenticated=True)
    
    async def get_balances(self) -> Dict[str, Decimal]:
        """Get account balances"""
        account_info = await self.get_account_info()
        balances = {}
        
        # Parse balances from account info
        if "assetPositions" in account_info:
            for asset in account_info["assetPositions"]:
                coin = asset.get("coin", "")
                total = Decimal(asset.get("total", "0"))
                balances[coin] = total
        
        return balances
    
    async def get_positions(self) -> List[Position]:
        """Get open positions"""
        account_info = await self.get_account_info()
        positions = []
        
        if "assetPositions" in account_info:
            for pos_data in account_info["assetPositions"]:
                if pos_data.get("position", {}).get("szi", "0") != "0":
                    position = Position(
                        symbol=pos_data.get("coin", ""),
                        side="long" if float(pos_data["position"]["szi"]) > 0 else "short",
                        size=Decimal(abs(float(pos_data["position"]["szi"]))),
                        entry_price=Decimal(pos_data["position"].get("entryPx", "0")),
                        mark_price=Decimal(pos_data.get("markPx", "0")),
                        unrealized_pnl=Decimal(pos_data["position"].get("unrealizedPnl", "0")),
                        margin=Decimal(pos_data["position"].get("marginUsed", "0"))
                    )
                    positions.append(position)
        
        return positions
    
    async def get_open_orders(self) -> List[Order]:
        """Get open orders"""
        data = await self._make_request("POST", "/info", 
                                      data={"type": "openOrders", "user": self.api_key},
                                      authenticated=True)
        
        orders = []
        for order_data in data:
            order = Order(
                order_id=str(order_data.get("oid", "")),
                symbol=order_data.get("coin", ""),
                side="buy" if order_data.get("side") == "B" else "sell",
                order_type="limit",  # Hyperliquid default
                amount=Decimal(order_data.get("sz", "0")),
                price=Decimal(order_data.get("limitPx", "0")),
                status="open",
                filled_amount=Decimal("0"),  # Would need additional API call
                timestamp=order_data.get("timestamp", int(time.time() * 1000))
            )
            orders.append(order)
        
        return orders
    
    # Trading Methods
    
    async def place_order(self, 
                         symbol: str,
                         side: str,
                         order_type: str,
                         amount: Decimal,
                         price: Optional[Decimal] = None,
                         time_in_force: str = "GTC",
                         reduce_only: bool = False) -> Dict[str, Any]:
        """Place an order"""
        
        order_data = {
            "coin": symbol,
            "is_buy": side.lower() == "buy",
            "sz": float(amount),
            "limit_px": float(price) if price else None,
            "order_type": {"limit": "limit", "market": "market"}.get(order_type.lower(), "limit"),
            "reduce_only": reduce_only
        }
        
        # Remove None values
        order_data = {k: v for k, v in order_data.items() if v is not None}
        
        return await self._make_request("POST", "/exchange", 
                                      data={"type": "order", "orders": [order_data]},
                                      authenticated=True)
    
    async def cancel_order(self, symbol: str, order_id: str) -> Dict[str, Any]:
        """Cancel an order"""
        return await self._make_request("POST", "/exchange",
                                      data={
                                          "type": "cancel",
                                          "cancels": [{
                                              "coin": symbol,
                                              "oid": int(order_id)
                                          }]
                                      },
                                      authenticated=True)
    
    async def cancel_all_orders(self, symbol: Optional[str] = None) -> Dict[str, Any]:
        """Cancel all orders for a symbol or all symbols"""
        data = {"type": "cancelByCloid"}
        if symbol:
            data["coin"] = symbol
        
        return await self._make_request("POST", "/exchange", data=data, authenticated=True)
    
    # WebSocket Methods
    
    async def connect_websocket(self):
        """Connect to WebSocket for real-time data"""
        try:
            self._ws = await websockets.connect(self.ws_url)
            self._logger.info("WebSocket connected")
        except Exception as e:
            self._logger.error(f"WebSocket connection failed: {e}")
            raise
    
    async def subscribe_orderbook(self, symbol: str):
        """Subscribe to order book updates"""
        if not self._ws:
            await self.connect_websocket()
        
        subscribe_msg = {
            "method": "subscribe",
            "subscription": {
                "type": "l2Book",
                "coin": symbol
            }
        }
        
        await self._ws.send(json.dumps(subscribe_msg))
        self._logger.info(f"Subscribed to {symbol} order book")
    
    async def subscribe_trades(self, symbol: str):
        """Subscribe to trade updates"""
        if not self._ws:
            await self.connect_websocket()
        
        subscribe_msg = {
            "method": "subscribe",
            "subscription": {
                "type": "trades",
                "coin": symbol
            }
        }
        
        await self._ws.send(json.dumps(subscribe_msg))
        self._logger.info(f"Subscribed to {symbol} trades")
    
    async def listen_websocket(self):
        """Listen for WebSocket messages"""
        if not self._ws:
            raise Exception("WebSocket not connected")
        
        try:
            async for message in self._ws:
                data = json.loads(message)
                yield data
        except websockets.exceptions.ConnectionClosed:
            self._logger.warning("WebSocket connection closed")
        except Exception as e:
            self._logger.error(f"WebSocket error: {e}")
            raise
    
    # Utility Methods
    
    async def get_server_time(self) -> int:
        """Get server timestamp"""
        return int(time.time() * 1000)
    
    async def ping(self) -> bool:
        """Test connectivity"""
        try:
            await self._make_request("GET", "/info")
            return True
        except Exception:
            return False
    
    def is_spot_trading_available(self) -> bool:
        """Check if spot trading is available (not in testnet)"""
        return not self.testnet